import torch
from torch.utils.data.sampler import BatchSampler
from easycore.common.registry import Registry

from .dataset import DatasetLoader
from .transform import TransformLoader
from .sampler import InfiniteSampler, InfiniteInstanceBalanceSampler, InfiniteCategoryBalanceSampler
from .data_mapper import MapDataset
from .collate import default_collate


class DataLoaderFactory:
    registry = Registry("data loader")

    @classmethod
    def register(cls, loader_type_name=None, obj=None):
        return cls.registry.register(loader_type_name, obj)

    @classmethod
    def build_train_loader(cls, loader_config, mapper=None):
        loader_config = {**loader_config['common'], **loader_config['train']}
        return cls.registry.get(loader_config['type']).build(loader_config, mapper)

    @classmethod
    def build_test_loader(cls, loader_config, mapper=None):
        loader_config = {**loader_config['common'], **loader_config['test']}
        return cls.register.get(loader_config['type']).build(loader_config, mapper)


@DataLoaderFactory.register()
class InfiniteLoader:
    @classmethod
    def build(cls, loader_config, mapper=None):
        batch_size = loader_config['batch_size']
        num_workers = loader_config['num_workers']

        dataset_config = loader_config['dataset']

        dataset = DatasetLoader.get(dataset_config)
        sampler = InfiniteSampler(len(dataset))
        sampler = BatchSampler(sampler, batch_size, drop_last=True)

        if mapper is not None:
            dataset = MapDataset(dataset, mapper)

        data_loader = torch.utils.data.DataLoader(
            dataset,
            batch_sampler=sampler,
            collate_fn=default_collate,
            num_workers=num_workers,
        )
        return data_loader


@DataLoaderFactory.register()
class InfiniteInstanceBalanceLoader:
    @classmethod
    def build(cls, loader_config, mapper=None):
        batch_size = loader_config['batch_size']
        num_workers = loader_config['num_workers']

        dataset_config = loader_config['dataset']

        dataset = DatasetLoader.get(dataset_config)
        sampler = InfiniteInstanceBalanceSampler(dataset, len(loader_config['categories']))
        sampler = BatchSampler(sampler, batch_size, drop_last=True)

        if mapper is not None:
            dataset = MapDataset(dataset, mapper)

        data_loader = torch.utils.data.DataLoader(
            dataset,
            batch_sampler=sampler,
            collate_fn=default_collate,
            num_workers=num_workers,
        )
        return data_loader


@DataLoaderFactory.register()
class InfiniteCategoryBalanceLoader:
    @classmethod
    def build(cls, loader_config, mapper=None):
        batch_size = loader_config['batch_size']
        num_workers = loader_config['num_workers']

        dataset_config = loader_config['dataset']

        dataset = DatasetLoader.get(dataset_config)
        sampler = InfiniteCategoryBalanceSampler(dataset, len(loader_config['categories']))
        sampler = BatchSampler(sampler, batch_size, drop_last=True)

        if mapper is not None:
            dataset = MapDataset(dataset, mapper)

        data_loader = torch.utils.data.DataLoader(
            dataset,
            batch_sampler=sampler,
            collate_fn=default_collate,
            num_workers=num_workers,
        )
        return data_loader
